package com.atguigu.study.controller;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

/**
 * @author Marco
 * @Description
 * @date 2025/8/22 2:20
 * @email 3293336923@qq.com
 */
@Slf4j
@RestController
public class EmbeddingController {

//    将 文档 向量化，形成 浮点型数组
    @Autowired
    private EmbeddingModel   embeddingModel;

    @Autowired
    private QdrantClient   qdrantClient;


//    将 浮点型 数组 存储在 向量数据库中
    @Autowired
    private EmbeddingStore<TextSegment>     embeddingStore;


    @GetMapping("/embedding/embed")
    public     String    embed(){
        String   prompt = """
                        床前明月光，
                        疑是地上霜，
                        举头望明月，
                        低头思故乡。
                """;

        Response<Embedding> response = embeddingModel.embed(prompt);
        log.info("response:{}",response);

        return  response.content().toString();
    }


    /***
     *  新建数据库实例 和 创建索引： test-qdrant
     *  类似 mysql  create  database  test-qdrant
     *  192.168.100.22:9012/embedding/createCollection
     */
    @GetMapping("/embedding/createCollection")
    public  void     createCollection(){
        Collections.VectorParams vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(384)
                .build();
        qdrantClient.createCollectionAsync("test-qdrant02" , vectorParams);
    }

    /**
     *  向 向量数据库中 新增文本记录
     * @return
     */
    @GetMapping("/embedding/add")
    public    String   add(){
        String   prompt = """
                    床前明月光，
                    疑是地上霜，
                    举头望明月，
                    低头思故乡。
                """;

        TextSegment segment = TextSegment.from(prompt);
        segment.metadata().put("author" , "marco");
        Embedding embedding = embeddingModel.embed(segment).content();

        String result = embeddingStore.add(embedding, segment);
//        返回一个流水号
// 42ae3d2a-0171-4444-a7b0-9df0b70525ffn

        return result;
    }


//  根据条件 去向量数据库中 进行 查找
    @GetMapping(value = "/embedding/query1")
    public void query1(){
        Embedding queryEmbedding = embeddingModel.embed("明月光").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);
        System.out.println(searchResult.matches().get(0).embedded().text());
    }



//    错误演示
    @GetMapping(value = "/embedding/query2")
    public void query2(){
        Embedding queryEmbedding = embeddingModel.embed("明月").content();

        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .filter(metadataKey("author").isEqualTo("marco1"))
                .maxResults(1)
                .build();

        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(embeddingSearchRequest);

        System.out.println(searchResult.matches().get(0).embedded().text());
    }
}
